package winter.com.ideaaedi.classwinter;

import winter.com.ideaaedi.classwinter.author.JustryDeng;
import winter.com.ideaaedi.classwinter.executor.DecryptExecutor;
import winter.com.ideaaedi.classwinter.util.Cache;
import winter.com.ideaaedi.classwinter.util.Constant;
import winter.com.ideaaedi.classwinter.util.ExceptionUtil;
import winter.com.ideaaedi.classwinter.util.IOUtil;
import winter.com.ideaaedi.classwinter.util.JarUtil;
import winter.com.ideaaedi.classwinter.util.JavaagentCmdArgs;
import winter.com.ideaaedi.classwinter.util.Logger;
import winter.com.ideaaedi.classwinter.util.Pair;
import winter.com.ideaaedi.classwinter.util.StrUtil;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.nio.charset.StandardCharsets;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.zip.ZipFile;

/**
 * 反向解密(Agent类)
 *
 * @author {@link JustryDeng}
 * @since 2021/5/6 22:38:36
 */
public class Reverses {
    
    private static JavaagentCmdArgs javaagentCmdArgs  = null;
    
    /**
     * pre-main 入口函数
     *
     * @param args
     *            参数
     * @param instrumentation
     *            This class provides services needed to instrumentation Java programming language code
     */
    public static void premain(String args, Instrumentation instrumentation) {
        if (javaagentCmdArgs == null) {
            javaagentCmdArgs = JavaagentCmdArgs.parseJavaagentCmdArgs(args);
            // 是否启动debug，同步给Logger
            Logger.ENABLE_DEBUG.set(javaagentCmdArgs.isDebug());
        }
        
        Logger.debug(Reverses.class, "Parse raw args [" + args + "] to javaagentCmdArgs -> " + javaagentCmdArgs);
    
        final AtomicBoolean firstExec = new AtomicBoolean(false);
    
        // 在JVM加载class字节码之前，通过ClassFileTransformer修改字节码
        if (instrumentation != null) {
            /*
             * 特别注意: 并不是说jar包中的所有class都会走到下面的逻辑中。
             *          只有jar包中被用到的class才会走到下面的逻辑中，不被使用的class是不会走到下面的逻辑的。
             *          也就是说: 如果加密时，加密了一个根本没有使用的class,那么该javaagent加载时，该class根本不会走到下面的逻辑中，进而不会走解密逻辑。
             */
            instrumentation.addTransformer(new ClassFileTransformer() {
                @Override
                public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
                    if (className == null || protectionDomain == null || loader == null) {
                        return classfileBuffer;
                    }
                    // 获取类所在的项目运行路径
                    String projectPath = protectionDomain.getCodeSource().getLocation().getPath();
                    projectPath = JarUtil.getRootPath(projectPath);
                    if (StrUtil.isEmpty(projectPath)) {
                        return classfileBuffer;
                    }
                    className = className.replace("/", ".").replace("\\", ".");
                    // 本项目的印章
                    if (Cache.sealCache == null) {
                        byte[] sealByte = IOUtil.readFileFromWorkbenchRoot(new File(projectPath), Constant.SEAL_FILE);
                        String sealContent = new String(sealByte, StandardCharsets.UTF_8);
                        if (StrUtil.isBlank(sealContent)) {
                            Logger.error(Reverses.class, "Obtain project seal fail.");
                            // 结束程序
                            System.exit(-1);
                        }
                        Logger.debug(Reverses.class, "seal of the project is -> " + sealContent);
                        Cache.sealCache = sealContent;
                    }
    
                    /// ============================================ 处理non-class文件
                    final String inputPwd = javaagentCmdArgs.getPassword();
                    if (firstExec.compareAndSet(false, true)) {
                        try {
                            // 解混淆 除了class文件外的其它文件
                            Map<String, Pair<byte[], byte[]>> resultMap =
                                    DecryptExecutor.unMaskNonClassFiles(projectPath, inputPwd == null ? null : inputPwd.toCharArray());
                            Map<String, byte[]> tmpMap = new HashMap<>(16);
                            resultMap.forEach((k, v) -> {
                                tmpMap.put(k, v.getLeft());
                            });
                            String finalProjectPath = projectPath;
                            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                                try {
                                    Logger.debug(Reverses.class, "mask non-class files start.");
                                    JarUtil.rewriteZipEntry(new ZipFile(finalProjectPath), tmpMap);
                                    Logger.debug(Reverses.class, "mask non-class files end.");
                                } catch (IOException e) {
                                    // ignore
                                }
                            }));
                        } catch (Exception e) {
                            Logger.error(Reverses.class, ExceptionUtil.getStackTraceMessage(e));
                            Logger.error(Reverses.class, "Decrypt non-classes fail. ");
                            System.exit(-1);
                            // 上一步System.exit(-1)就退出程序了，照理说是不会走到下面这里的(，不过为了以防万一，这里打出提醒, class-winter失效)
                            for (int i = 0; i < Constant.TEN; i++) {
                                Logger.error(Reverses.class, "!!!!!!!!!!! class-winter Invalidation. !!!!!!!!!!!");
                            }
                            return null;
                        }
                    }
    
                    /// ============================================ 处理class文件
                    // 判断是否应该解密
                    if (DecryptExecutor.checklistContain(projectPath, className) && DecryptExecutor.verifySeal(classfileBuffer)) {
                        //noinspection DuplicatedCode
                        try {
                            Logger.debug(Reverses.class, "Decrypt class[" + className + "] start.");
                            classfileBuffer = DecryptExecutor.process(projectPath, null, className,
                                    inputPwd == null ? null :
                                    inputPwd.toCharArray());
                            Logger.debug(Reverses.class, "Decrypt class[" + className + "] end.");
                            return classfileBuffer;
                        } catch (Exception e) {
                            Logger.error(Reverses.class, ExceptionUtil.getStackTraceMessage(e));
                            Logger.error(Reverses.class, "Decrypt class[" + className + "] fail. "
                                    + (StrUtil.isEmpty(inputPwd) ? e.getMessage() : "Please ensure your password is correct."));
                            System.exit(-1);
                            // 上一步System.exit(-1)就退出程序了，照理说是不会走到下面这里的(，不过为了以防万一，这里打出提醒, class-winter失效)
                            for (int i = 0; i < Constant.TEN; i++) {
                                Logger.error(Reverses.class, "!!!!!!!!!!! class-winter Invalidation. !!!!!!!!!!!");
                            }
                            return null;
                        }
                    } else if (DecryptExecutor.checklistOfAllLibsContain(projectPath, className)
                            && DecryptExecutor.verifyLibSeal(projectPath, className, classfileBuffer)) {
                        String classWinterInfoDir = DecryptExecutor.getLibDirRelativePath(className);
                        //noinspection DuplicatedCode
                        try {
                            Logger.debug(Reverses.class, "Decrypt class[" + className + "] start.");
                            // lib中的密码在加密时，都统一处理好了，这里直接传null
                            classfileBuffer = DecryptExecutor.process(projectPath, classWinterInfoDir, className, null);
                            Logger.debug(Reverses.class, "Decrypt class[" + className + "] end.");
                            return classfileBuffer;
                        } catch (Exception e) {
                            Logger.error(Reverses.class, ExceptionUtil.getStackTraceMessage(e));
                            String lib = DecryptExecutor.parseLib(classWinterInfoDir);
                            Logger.error(Reverses.class, "Decrypt class[" + className + "] fail. \nPlease check:\n"
                                    + "\t1. Ensure 'Your lib " + lib + " need a input password ?'\n"
                                    + "\t2. Ensure 'Your lib "+ lib +"'s password is correct ?'");
                            System.exit(-1);
                            // 上一步System.exit(-1)就退出程序了，照理说是不会走到下面这里的(，不过为了以防万一，这里打出提醒, class-winter失效)
                            for (int i = 0; i < Constant.TEN; i++) {
                                Logger.error(Reverses.class, "!!!!!!!!!!! class-winter Invalidation. !!!!!!!!!!!");
                            }
                            return null;
                        }
                    } else {
                        return classfileBuffer;
                    }
                }
            });
        }
    }
    

}
